//Author: 8891689
//https://github.com/8891689
#ifndef AES_H
#define AES_H

#include <stdint.h>
#include <stddef.h> // For size_t

// 定义 AES 密钥长度类型
typedef enum {
    AES_KEY_128 = 4, // 128 bit = 4 * 32-bit words
    AES_KEY_192 = 6, // 192 bit = 6 * 32-bit words
    AES_KEY_256 = 8  // 256 bit = 8 * 32-bit words
} AES_KeySize;

// 定义轮密钥生成模式
typedef enum {
    AES_KEY_SCHEDULE_ENCRYPT, // 为加密生成轮密钥
    AES_KEY_SCHEDULE_DECRYPT  // 为解密生成优化轮密钥
} AES_KeySchedule_Mode;


#define AES_BLOCK_SIZE 16 // 区块大小为 16 bytes
#define AES_Nb 4          // 状态矩阵的列数 (4 * 32-bit words)
#define AES_Nr(Nk_enum_val) (6 + (Nk_enum_val)) // 根据密钥长度计算回合数


/**
 * @brief 生成 AES 轮密钥。
 * @param key 输入的原始密钥。
 * @param w 输出轮密钥的数组。
 * @param Nk 密钥长度类型 (AES_KEY_128, AES_KEY_192, AES_KEY_256)。
 * @param mode 指定为加密或解密生成轮密钥。
 * @return 成功返回 0。
 */
int AES_KeyExpansion(const uint8_t *key, uint32_t *w, AES_KeySize Nk, AES_KeySchedule_Mode mode);

/**
 * @brief 使用 T-Table 优化方法加密一个 16-byte 的区块。
 * @param in_out 指向 16-byte 数据的指针，加密后结果会写回此处。
 * @param w_encrypt 加密用的轮密钥 (由 AES_KeyExpansion 生成，模式为 AES_KEY_SCHEDULE_ENCRYPT)。
 * @param Nk 密钥长度类型。
 */
void AES_Encrypt(uint8_t *in_out, const uint32_t *w_encrypt, AES_KeySize Nk);

/**
 * @brief 使用 T-Table 优化方法解密一个 16-byte 的区块。
 * @param in_out 指向 16-byte 数据的指针，解密后结果会写回此处。
 * @param w_decrypt 解密用的轮密钥 (由 AES_KeyExpansion 生成，模式为 AES_KEY_SCHEDULE_DECRYPT)。
 * @param Nk 密钥长度类型。
 */
void AES_Decrypt(uint8_t *in_out, const uint32_t *w_decrypt, AES_KeySize Nk);

/**
 * @brief 打印 16-byte 的状态数组 (用于调试)。
 * @param state 指向 16-byte 数据的指针。
 */
void AES_PrintState(const uint8_t *state);

#endif // AES_H
